import numpy as np
import unittest
import os
import sys
import time
import os.path
import glob
import logging
import random
import re
import datetime
import subprocess
import json

import npu_device
import tensorflow.compat.v1 as tf
from npu_device.compat.v1.estimator import npu_ops
from tensorflow.core.protobuf.rewriter_config_pb2 import RewriterConfig
npu_device.compat.enable_v1()
tf.disable_v2_behavior()

# os.environ["ASCEND_OP_COMPILER_WORK_PATH_IN_KERNEL_META"] = '1'
atol = 0.01
rtol = 0.01
log_path = './AFlog/'
# OFFLINE = False # True: NPU run， false: cpu run
# OFFLINE = True # True: NPU run， false: cpu run
min = -10
middle = 0
max = 10

is_npu = 1


class TestAutoFuseCase(unittest.TestCase):
    def setUp(self):  # All case before process
        os.system('rm  ge*txt')
        os.system('mkdir %s' % log_path)
        os.environ['EXPERIMENTAL_ENABLE_AUTOFUSE'] = '1'
        os.environ['EXPERIMENTAL_LOWERING_CONCAT'] = '1'
        
        if len(sys.argv) >= 3:  #modify Tensor format
            self.input_format = sys.argv[-2]
            self.input_shape = list(map(int, sys.argv[-3].split(",")))
            self.case_name = sys.argv[-1]
        else:
            dtype = random.choice([0, 1])  # random choise dtype， 变量名字叫错input_format
            if dtype == 0:
                self.input_format = np.float32
            elif dtype == 1:
                self.input_format = np.float16
            else:
                self.input_format = np.int32
            self.input_shape = [random.choice([random.randint(1, 20)]) for x in
                                range(random.randint(1, 4))]  # random choise shape
            self.case_name = sys.argv[-1]
        self.str_format = str(self.input_format).split('.')[1].split("'")[0]
        
        self.log = logging
        LOG_FORMAT = "%(asctime)s [%(levelname)s]:%(message)s"
        time_stamp = datetime.datetime.now()
        logfile = log_path + self.case_name + '.log'
        self.log.basicConfig(filename=logfile, level=self.log.DEBUG, format=LOG_FORMAT, filemode='w')
        tf.compat.v1.reset_default_graph()  # clear graph
        
    def tearDown(self):  # case after process
        pass
    
    def build_session_config(self, execution_type):
        if execution_type == 'NPU':
            self.session_config = tf.compat.v1.ConfigProto()
            self.session_config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
            self.custom_op = self.session_config.graph_options.rewrite_options.custom_optimizers.add()
            self.custom_op.name = "NpuOptimizer"
            self.custom_op.parameter_map["enable_data_pre_proc"].b = True
            self.custom_op.parameter_map["mix_compile_mode"].b = False
            self.custom_op.parameter_map["use_off_line"].b = True
            self.custom_op.parameter_map["min_group_size"].b = 1
            self.custom_op.parameter_map["precision_mode"].s = tf.compat.as_bytes("allow_mix_precision")
            
        elif execution_type == 'CPU':
            self.session_config = tf.compat.v1.ConfigProto(
                allow_soft_placement=True,
                log_device_placement=False)
        self.session_config.gpu_options.allow_growth = True
        
    def sess_run(self, output, feed_dict):
        self.build_session_config("CPU")
        with tf.compat.v1.Session(config=self.session_config) as sess:
            result_cpu = sess.run(output, feed_dict=feed_dict)
            print("end execute CPU")
            
        self.build_session_config("NPU")
        with tf.compat.v1.Session(config=self.session_config) as sess:
            result_npu = sess.run(output, feed_dict=feed_dict)
            print("end execute NPU")
            
            
            self.assertTrue(np.allclose(result_npu, result_cpu, atol, rtol),
                            "AF-smoke-case-fail")  # exp_result
            print("AF-smoke-case-pass")
            print('end execute DaVinci')
            
    def test_tc_af_tf1_can_fuse_verticalFuse_0002(self):
        case_name = 'test_tc_af_tf1_can_fuse_verticalFuse_0002'
        try:
            self.input_dtype = np.float32  # 写死dtype
            self.input_shape = [64, 64, 64, 64]
            self.input_shape2 = [64, 64, 64]
            print('begin execute TestCase {}'.format(case_name))
            
            data1 = np.random.uniform(min, max, self.input_shape).astype(self.input_dtype)
            d1 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape)
            
            data2 = np.random.uniform(min, max, self.input_shape2).astype(self.input_dtype)
            d2 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape2)
            feed_dict = {d1: data1, d2: data2}
            
            print("prepare build graph")
            abs_1 = tf.abs(d1)
            abs_2 = tf.abs(abs_1)
            abs_3 = tf.abs(abs_2)
            
            with tf.control_dependencies([abs_3]):
                relu_1 = tf.nn.relu(d2)
                abs_4 = tf.abs(relu_1)
                
            add_1 = tf.add(abs_3, abs_4)
            output = add_1
            
            self.sess_run(output, feed_dict)
            
        except ValueError as e:
            print(e)
        finally:
            print('TestCase {} finished'.format(case_name))
            
    def test_tc_af_tf1_can_fuse_horizontalFuse_0004(self):
        case_name = 'test_tc_af_tf1_can_fuse_horizontalFuse_0004'
        try:
            self.input_dtype = np.float32  # 写死dtype
            self.input_shape = [14,20,20,20]
            self.input_shape2 = [6,14,20,20,20]
            print('begin execute TestCase {}'.format(case_name))
            
            data1 = np.random.uniform(min, max, self.input_shape).astype(self.input_dtype)
            d1 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape)
            
            data2 = np.random.uniform(min, max, self.input_shape2).astype(self.input_dtype)
            d2 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape2)
            feed_dict = {d1: data1, d2: data2}
            
            print("prepare build graph")
            abs_1 = tf.abs(d1)
            abs_2 = tf.abs(abs_1)
            abs_3 = tf.abs(abs_2)
            const_0 = tf.constant(1.05, shape = self.input_shape, dtype=self.input_dtype)
            matmul_0 = tf.matmul(abs_3, const_0)
            relu_1 = tf.nn.relu(matmul_0)
            add_1 = tf.add(matmul_0, d2)
            add_2 = tf.add(relu_1, add_1)
            output = add_2
            
            self.sess_run(output, feed_dict)
            
        except ValueError as e:
            print(e)
        finally:
            print('TestCase {} finished'.format(case_name))
            
    def test_tc_af_tf1_schedule_gen_tilinggroup_0001(self):
        case_name = 'test_tc_af_tf1_schedule_gen_tilinggroup_0001'
        try:
            self.input_dtype = np.float32  # 写死dtype
            self.input_shape = [1024, 48]
            self.input_shape2 = [1024, 1]
            print('begin execute TestCase {}'.format(case_name))
            
            data1 = np.random.uniform(min, max, self.input_shape).astype(self.input_dtype)
            d1 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape)
            
            data2 = np.random.uniform(min, max, self.input_shape2).astype(self.input_dtype)
            d2 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape2)
            
            print("prepare build graph")
            
            abs_1 = tf.math.abs(d1)
            relu_1 = tf.nn.relu(abs_1)
            add_1 = tf.math.add(relu_1, d2)
            
            reduce_sum_1 = tf.reduce_sum(add_1, axis=-1)
            square_1= tf.math.square(reduce_sum_1)
            output = square_1
            
            feed_dict = {d1: data1, d2: data2}
            
            self.sess_run(output, feed_dict)
            
        except ValueError as e:
            print(e)
        finally:
            print('TestCase {} finished'.format(case_name))
            
    def test_tc_af_tf1_schedule_gen_tilinggroup_0002(self):
        case_name = 'test_tc_af_tf1_schedule_gen_tilinggroup_0002'
        try:
            self.input_dtype = np.float32  # 写死dtype
            self.input_shape = [16, 32, 64, 128]
            self.input_shape2 = [16, 32, 64, 1]
            self.input_shape3 = [16, 32, 64]
            print('begin execute TestCase {}'.format(case_name))
            
            data1 = np.random.uniform(min, max, self.input_shape).astype(self.input_dtype)
            d1 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape)
            
            data2 = np.random.uniform(min, max, self.input_shape2).astype(self.input_dtype)
            d2 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape2)
            
            data3 = np.random.uniform(min, max, self.input_shape3).astype(self.input_dtype)
            d3 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape3)
            
            print("prepare build graph")
            
            abs_1 = tf.math.exp(d1)
            relu_1 = tf.nn.relu(abs_1)
            add_1 = tf.math.multiply(relu_1, d2)
            
            reduce_sum_1 = tf.reduce_sum(add_1, axis=-1)
            square_1= tf.math.square(reduce_sum_1)
            concat_1 = tf.concat([square_1, d3], axis=-1)
            output = concat_1
            
            feed_dict = {d1: data1, d2: data2, d3: data3}
            
            self.sess_run(output, feed_dict)
            
        except ValueError as e:
            print(e)
        finally:
            print('TestCase {} finished'.format(case_name))
            
    def test_tc_af_tf1_schedule_gen_tilinggroup_0019(self):
        case_name = 'test_tc_af_tf1_schedule_gen_tilinggroup_0019'
        try:
            self.input_dtype = np.float32  # 写死dtype
            self.input_shape = [1024, 32]
            self.input_shape2 = [1024, 1]
            self.input_shape3 = [1024,]
            print('begin execute TestCase {}'.format(case_name))
            
            data1 = np.random.uniform(min, max, self.input_shape).astype(self.input_dtype)
            d1 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape)
            
            data2 = np.random.uniform(min, max, self.input_shape2).astype(self.input_dtype)
            d2 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape2)
            
            data3 = np.random.uniform(min, max, self.input_shape3).astype(self.input_dtype)
            d3 = tf.compat.v1.placeholder(self.input_dtype, shape=self.input_shape3)
            
            print("prepare build graph")
            
            abs_1 = tf.math.abs(data1)
            relu_1 = tf.nn.relu(abs_1)
            add_1 = tf.math.add(relu_1, data2)
            reduce_sum_1 = tf.reduce_sum(add_1, axis=-1)
            square_1= tf.math.square(reduce_sum_1)
            output = tf.concat([square_1, data3], axis=0)
            feed_dict = {d1: data1, d2: data2, d3: data3}
            
            self.sess_run(output, feed_dict)
            
        except ValueError as e:
            print(e)
        finally:
            print('TestCase {} finished'.format(case_name))

if __name__ == '__main__':
    suite = unittest.TestSuite()
    if len(sys.argv) >= 2:
        suite.addTest(TestAutoFuseCase(sys.argv[-1]))
    else:
        suite.addTest(TestAutoFuseCase('test_tc_af_tf1_can_fuse_verticalFuse_0002'))
    unittest.TextTestRunner().run(suite)
